from torch.utils.data import Dataset,DataLoader
import os
from PIL import Image
from my_transform import img_transforms

root = os.path.join("../..","data/VOC2012/")
#去读取VOC2012数据集在语义分割上的使用的图片的名称,当train为True时，使用的是训练数据
def read_image(root = root,is_train=True):
    root_train = root + 'ImageSets/Segmentation/' + ('train.txt' if is_train else 'val.txt')

    with open(root_train,'r') as f:
        images_name = f.read().split()#读取所有的数据,所以要清理一下空白字符，来划分

    print("总共读取了{0}张图片".format(len(images_name)))
    data = [os.path.join(root,'JPEGImages',i+'.jpg') for i in images_name]
    label= [os.path.join(root,'SegmentationClass',i+'.png') for i in images_name]

    return data,label








# print(root)
# read_image(root)

class myDataset(Dataset):

    def __init__(self,root,train,crop_size,trainform=None):

        self.trainform = trainform
        self.crop_szie = crop_size
        self.root = root
        data_list,label_list  = read_image(root,is_train=train)

        self.data_list = self.filter(data_list)
        self.label_list= self.filter(label_list)

        print("总共读取了{0}张图片原件和{1}".format(len(self.data_list),len(self.label_list)))



    def filter(self,images_name):
        return [img_name for img_name in images_name if(Image.open(img_name).size[0] >= self.crop_szie[0] and
                                                        Image.open(img_name).size[1] >= self.crop_szie[1])]


    def __len__(self):
        return len(self.data_list)

    def __getitem__(self, item):
        data_name = self.data_list[item]
        label_name = self.label_list[item]

        img_data,img_label = self.trainform(data_name,label_name,self.crop_size)

        return img_data,img_label
        img,label  = self.trainform(data_name,label_name,self.crop_szie)

        return img,label#返回两个图片一个正常图片，一个标注过的图片








def my_train_test():
    trainbag = myDataset(train=True,trainform = img_transforms)
    train_data = DataLoader(trainbag,batch_size=4,shuffle=True,num_workers=4)

    testbag  = myDataset(train=False,trainform= img_transforms)
    test_data = DataLoader(testbag,batch_size=4,shuffle=True,num_workers=4)


    return train_data,test_data


